package com.naiterui.ehp.bs.file.api.filter;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.nio.charset.StandardCharsets;

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;

/**
 * @author gaoliangliang
 * @since 2021/9/16
 */
public class ResponseWrapper extends HttpServletResponseWrapper {

    private final ByteArrayOutputStream bytes = new ByteArrayOutputStream();
    private PrintWriter pwrite;

    public ResponseWrapper(HttpServletResponse response) {
        super(response);
    }

    @Override
    public ServletOutputStream getOutputStream() throws IOException {
        // 将数据写到 byte 中
        return new CacheServletOutputStream(bytes);
    }

    /**
     * 重写父类的 getWriter() 方法，将响应数据缓存在 PrintWriter 中
     */
    @Override
    public PrintWriter getWriter() throws IOException {
        pwrite = new PrintWriter(new OutputStreamWriter(bytes, StandardCharsets.UTF_8));
        return pwrite;
    }

    /**
     * PrintWriter 缓存响应数据
     * 
     * @return 响应字节数组
     */
    public byte[] getBytes() {
        if (null != pwrite) {
            pwrite.close();
            return bytes.toByteArray();
        }

        try {
            bytes.flush();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return bytes.toByteArray();
    }

    static class CacheServletOutputStream extends ServletOutputStream {
        private final ByteArrayOutputStream cacheStream;

        public CacheServletOutputStream(ByteArrayOutputStream cacheStream) {
            this.cacheStream = cacheStream;
        }

        @Override
        public void write(int b) throws IOException {
            // 将数据写到 cacheStream 中
            cacheStream.write(b);
        }

        @Override
        public boolean isReady() {
            return true;
        }

        @Override
        public void setWriteListener(WriteListener listener) {
            throw new UnsupportedOperationException();
        }
    }
}